import os
from PIL import Image
import numpy as np


def compute_mean_std():
    img_channels = 1
    img_dir = "SLIVER07/training/images"
    roi_dir = "SLIVER07/training/mask"
    assert os.path.exists(img_dir), f"image dir: '{img_dir}' does not exist."
    assert os.path.exists(roi_dir), f"roi dir: '{roi_dir}' does not exist."

    img_name_list = [i for i in os.listdir(img_dir) if i.endswith(".png")]
    cumulative_mean = np.zeros(img_channels)
    cumulative_std = np.zeros(img_channels)
    for img_name in img_name_list:
        img_path = os.path.join(img_dir, img_name)
        ori_path = os.path.join(roi_dir, img_name)
        img = np.array(Image.open(img_path).convert('L')) / 255.
        roi_img = np.array(Image.open(ori_path).convert('L'))

        img = img[roi_img == 255]                                                                        #用于从图像 img 中提取位于特定区域的像素值。
        cumulative_mean += img.mean(axis=0)                                                     #axis=0: 表示沿着数组的第一个轴（即行轴）进行操作，这里指沿着行的方向计算平均值。在二维数组中，沿着行的方向对每列进行操作，即对每个像素列的值求平均。
        cumulative_std += img.std(axis=0)

    mean = cumulative_mean / len(img_name_list)
    mean_result = np.round(mean, decimals=2)[0]
    std = cumulative_std / len(img_name_list)
    std_result = np.round(std, decimals=2)[0]
    return mean_result,std_result


if __name__ == '__main__':
    mean,std = compute_mean_std()
    print(f"mean: {mean}")
    print(f"std: {std}")
